#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
CTP行情接收并合成1分钟K线记录程序
接收6个品种的行情数据，使用正确的K线合成算法，保存到CSV文件
按照wtpy标准：1分钟K线应该是60秒内tick数据的累积统计
"""

import ai_ctp
import json
import threading
import time
import csv
from pathlib import Path
from datetime import datetime, timedelta
from typing import Dict, List, Optional
import copy

class TickData:
    """Tick数据存储类"""
    def __init__(self):
        self.instrument_id = ""
        self.last_price = 0.0
        self.volume = 0
        self.turnover = 0.0
        self.open_interest = 0
        self.update_time = ""
        self.action_day = ""
        self.timestamp = 0  # 添加时间戳字段

class KlineBar:
    """K线数据存储类 - 参考wtpy的Bar结构"""
    def __init__(self):
        self.instrument_id = ""
        self.bar_time = 0           # K线时间戳
        self.time_str = ""          # 时间字符串
        self.period = 1             # 周期（分钟）
        self.open = 0.0             # 开盘价（第一个tick）
        self.high = 0.0             # 最高价（期间内最高）
        self.low = 0.0              # 最低价（期间内最低）
        self.close = 0.0            # 收盘价（最后一个tick）
        self.volume = 0             # 成交量（期间内累计增量）
        self.turnover = 0.0         # 成交额（期间内累计增量）
        self.open_interest = 0      # 持仓量（最新值）
        self.tick_count = 0         # tick数量统计
        self.first_tick_time = ""   # 第一个tick时间
        self.last_tick_time = ""    # 最后一个tick时间
        self.is_completed = False   # 是否已完成

class KlineRecorderCallback(ai_ctp.PyMdCallback):
    """K线记录回调类 - 按照wtpy标准实现正确的K线合成"""
    
    def __init__(self, data_dir="./kline_data"):
        super().__init__()
        
        # 基础配置
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(exist_ok=True)
        
        # 线程安全锁
        self.lock = threading.Lock()
        
        # 订阅的合约列表
        self.instruments = ["au2508", "IC2509", "IF2509", "IH2509", "IM2509", "rb2510"]
        
        # 统计信息
        self.tick_count = 0
        self.kline_count = 0
        
        # K线数据缓存 - 每个合约的当前K线
        self.current_klines: Dict[str, Optional[KlineBar]] = {}
        
        # tick数据缓存 - 用于计算增量
        self.previous_ticks: Dict[str, Optional[TickData]] = {}
        
        # 初始化缓存
        for instrument in self.instruments:
            self.current_klines[instrument] = None
            self.previous_ticks[instrument] = None
            
            # 创建数据目录
            instrument_dir = self.data_dir / instrument
            instrument_dir.mkdir(exist_ok=True)
        
        print(f"✓ 初始化K线记录器，数据目录: {self.data_dir}")
        print(f"✓ 订阅合约: {', '.join(self.instruments)}")
        print("📈 使用wtpy标准K线合成算法 - 时间段内累计统计")

    def on_connected(self):
        print("✓ CTP行情服务器连接成功")

    def on_disconnected(self, reason):
        print(f"✗ CTP行情服务器断开连接: {reason}")

    def on_login_success(self):
        print("✓ CTP行情登录成功")

    def on_login_failed(self, error_id, error_msg):
        print(f"✗ CTP行情登录失败: {error_id} - {error_msg}")

    def on_tick(self, tick):
        """处理tick数据 - 按照wtpy标准实现正确的K线合成"""
        with self.lock:
            self.tick_count += 1
            
            # 转换tick数据
            tick_data = self._convert_tick_data(tick)
            if tick_data is None:
                return
            
            # 执行K线合成
            self._synthesize_kline(tick_data)
            
            if self.tick_count % 50 == 0:  # 每50个tick打印一次
                print(f"Tick #{self.tick_count}: {tick.instrument_id} 价格:{tick.last_price:.3f} "
                      f"成交量:{tick.volume} 时间:{tick.update_time}")

    def _convert_tick_data(self, tick) -> Optional[TickData]:
        """转换tick数据格式"""
        try:
            tick_data = TickData()
            tick_data.instrument_id = tick.instrument_id
            tick_data.last_price = tick.last_price
            tick_data.volume = tick.volume
            tick_data.turnover = tick.turnover
            tick_data.open_interest = tick.open_interest
            tick_data.update_time = tick.update_time
            tick_data.action_day = tick.action_day
            
            # 计算时间戳
            tick_data.timestamp = self._calculate_timestamp(tick.update_time)
            
            return tick_data
        except Exception as e:
            print(f"✗ 转换tick数据失败: {e}")
            return None

    def _calculate_timestamp(self, time_str: str) -> int:
        """计算tick时间戳"""
        try:
            hour, minute, second = map(int, time_str.split(':'))
            now = datetime.now()
            tick_time = now.replace(hour=hour, minute=minute, second=second, microsecond=0)
            return int(tick_time.timestamp())
        except:
            return int(datetime.now().timestamp())

    def _synthesize_kline(self, tick_data: TickData):
        """K线合成 - 按照wtpy标准实现时间段内累计统计"""
        instrument_id = tick_data.instrument_id
        if instrument_id not in self.instruments:
            return
            
        # 计算1分钟K线时间戳（左闭右开区间）
        bar_time = self._calculate_bar_time(tick_data.timestamp, 1)
        if bar_time == 0:
            return
            
        current_kline = self.current_klines[instrument_id]
        previous_tick = self.previous_ticks[instrument_id]
        
        # 判断是否需要创建新K线
        is_new_bar = (current_kline is None or current_kline.bar_time != bar_time)
        
        if is_new_bar:
            # 保存上一根完成的K线
            if current_kline is not None:
                current_kline.is_completed = True
                current_kline.last_tick_time = previous_tick.update_time if previous_tick else ""
                self._save_completed_kline(current_kline)
                
            # 创建新K线 - 开盘价为第一个tick价格
            new_kline = KlineBar()
            new_kline.instrument_id = instrument_id
            new_kline.bar_time = bar_time
            new_kline.time_str = self._format_bar_time(bar_time)
            new_kline.period = 1
            new_kline.open = tick_data.last_price      # 开盘价：第一个tick
            new_kline.high = tick_data.last_price      # 初始最高价
            new_kline.low = tick_data.last_price       # 初始最低价
            new_kline.close = tick_data.last_price     # 收盘价：当前tick（会持续更新）
            new_kline.open_interest = tick_data.open_interest
            new_kline.tick_count = 1
            new_kline.first_tick_time = tick_data.update_time
            new_kline.last_tick_time = tick_data.update_time
            
            # 计算成交量和成交额增量
            if previous_tick is not None:
                # 基于前一个tick计算增量
                volume_increment = max(0, tick_data.volume - previous_tick.volume)
                turnover_increment = max(0, tick_data.turnover - previous_tick.turnover)
                new_kline.volume = volume_increment
                new_kline.turnover = turnover_increment
            else:
                # 第一个tick，直接使用当前值
                new_kline.volume = tick_data.volume
                new_kline.turnover = tick_data.turnover
            
            self.current_klines[instrument_id] = new_kline
            
            print(f"✓ 新建K线: {instrument_id} {new_kline.time_str} "
                  f"开盘:{new_kline.open:.3f} 成交量:{new_kline.volume} "
                  f"[时间段:{new_kline.first_tick_time}-]")
        else:
            # 更新现有K线 - 累计统计
            current_kline.high = max(current_kline.high, tick_data.last_price)  # 更新最高价
            current_kline.low = min(current_kline.low, tick_data.last_price)    # 更新最低价
            current_kline.close = tick_data.last_price                          # 收盘价：最新tick
            current_kline.open_interest = tick_data.open_interest               # 持仓量：最新值
            current_kline.tick_count += 1
            current_kline.last_tick_time = tick_data.update_time
            
            # 累计成交量和成交额增量
            if previous_tick is not None:
                volume_increment = max(0, tick_data.volume - previous_tick.volume)
                turnover_increment = max(0, tick_data.turnover - previous_tick.turnover)
                current_kline.volume += volume_increment
                current_kline.turnover += turnover_increment
            
            # 定期打印更新信息
            if current_kline.tick_count % 20 == 0:
                print(f"✓ 更新K线: {instrument_id} tick#{current_kline.tick_count} "
                      f"OHLC({current_kline.open:.3f},{current_kline.high:.3f},"
                      f"{current_kline.low:.3f},{current_kline.close:.3f}) "
                      f"累计量:{current_kline.volume} "
                      f"[{current_kline.first_tick_time}-{current_kline.last_tick_time}]")
        
        # 更新前一个tick缓存
        self.previous_ticks[instrument_id] = tick_data

    def _calculate_bar_time(self, timestamp: int, period_minutes: int) -> int:
        """
        计算K线时间戳 - 按照期货行业标准
        
        规则：左闭右开区间 [13:29:00, 13:30:00)
        - 13:29:00-13:29:59的tick → 13:29分钟K线
        - 13:30:00开始 → 13:30分钟K线
        
        返回值：K线标记时间戳
        """
        try:
            dt = datetime.fromtimestamp(timestamp)
            
            # 计算K线所属的分钟数（左闭右开区间）
            bar_minute = (dt.minute // period_minutes) * period_minutes
            
            # 构造K线时间戳
            bar_datetime = dt.replace(minute=bar_minute, second=0, microsecond=0)
            
            return int(bar_datetime.timestamp())
            
        except Exception as e:
            print(f"✗ 计算K线时间戳失败: {timestamp}, 错误: {e}")
            return 0

    def _format_bar_time(self, bar_time: int) -> str:
        """格式化K线时间"""
        try:
            dt = datetime.fromtimestamp(bar_time)
            return dt.strftime("%Y%m%d %H:%M:%S")
        except:
            return str(bar_time)

    def _save_completed_kline(self, kline: KlineBar):
        """保存完成的K线数据到CSV"""
        try:
            self.kline_count += 1
            
            print(f"📊 完成K线 #{self.kline_count}: {kline.instrument_id} 1分钟 "
                  f"时间:{kline.time_str} "
                  f"OHLC:({kline.open:.3f},{kline.high:.3f},"
                  f"{kline.low:.3f},{kline.close:.3f}) "
                  f"成交量:{kline.volume} Tick数:{kline.tick_count} "
                  f"时段:[{kline.first_tick_time}-{kline.last_tick_time}] [已闭线]")
            
            # 保存到CSV文件
            self._save_kline_to_csv(kline)
            
        except Exception as e:
            print(f"✗ 保存完成K线失败: {e}")

    def on_kline(self, kline):
        """接收到接口K线数据的回调（作为备用验证）"""
        try:
            if kline.period == 1:
                print(f"📊 [接口K线] {kline.instrument_id} {kline.time_str} "
                      f"OHLC({kline.open:.3f},{kline.high:.3f},{kline.low:.3f},{kline.close:.3f}) "
                      f"成交量:{kline.volume} (接口合成，仅供对比)")
        except Exception as e:
            print(f"⚠ K线回调处理错误: {e}")

    def _save_kline_to_csv(self, kline: KlineBar):
        """保存K线数据到CSV文件"""
        try:
            # 获取当前日期作为文件名
            today = datetime.now().strftime("%Y%m%d")
            instrument_dir = self.data_dir / kline.instrument_id
            csv_file = instrument_dir / f"{kline.instrument_id}_1min_{today}.csv"
            
            # 检查文件是否存在，不存在则创建并写入标题行
            file_exists = csv_file.exists()
            
            with open(csv_file, 'a', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                
                # 如果是新文件，写入标题行
                if not file_exists:
                    writer.writerow([
                        'instrument_id', 'datetime', 'bar_time', 'period',
                        'open', 'high', 'low', 'close', 
                        'volume', 'turnover', 'open_interest', 'tick_count', 
                        'first_tick_time', 'last_tick_time', 'is_completed'
                    ])
                
                # 写入K线数据
                writer.writerow([
                    kline.instrument_id,
                    kline.time_str,
                    kline.bar_time,
                    kline.period,
                    kline.open,
                    kline.high,
                    kline.low,
                    kline.close,
                    kline.volume,
                    kline.turnover,
                    kline.open_interest,
                    kline.tick_count,
                    kline.first_tick_time,
                    kline.last_tick_time,
                    kline.is_completed
                ])
                
        except Exception as e:
            print(f"✗ 保存K线数据失败: {e}")

    def on_error(self, error_id, error_msg):
        print(f"✗ CTP错误: {error_id} - {error_msg}")

    def get_status(self):
        """获取运行状态"""
        with self.lock:
            return {
                'tick_count': self.tick_count,
                'kline_count': self.kline_count,
                'instruments': len(self.instruments)
            }

class KlineRecorder:
    """K线记录器主类"""
    
    def __init__(self, config_file="ctp_simnow.json"):
        self.config_file = config_file
        self.config = {}
        self.md = None
        self.callback = None
        self.running = False
        
        # 加载配置
        self._load_config()
        
        print("✓ K线记录器初始化完成")
        print("📈 使用wtpy标准K线合成 - 确保时间段内累计统计的准确性")

    def _load_config(self):
        """加载CTP配置"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                self.config = json.load(f)
            print(f"✓ 加载配置文件: {self.config_file}")
        except Exception as e:
            print(f"✗ 加载配置文件失败: {e}")
            raise

    def start(self):
        """启动K线记录器"""
        print("启动专业级K线记录器...")
        print("📈 基于wtpy标准的K线合成算法")
        
        try:
            # 创建回调对象
            self.callback = KlineRecorderCallback()
            
            # 创建行情接口
            self.md = ai_ctp.MdApi()
            self.md.set_callback(self.callback)
            
            # 从配置文件的md部分读取配置
            md_config = self.config['md']
            
            # 使用AI_CTP的正确初始化方法
            result = self.md.init(
                md_config['broker_id'],
                md_config['user_id'], 
                md_config['password'],
                md_config['front_addresses'][0]  # 使用第一个前置地址
            )
            
            if not result:
                raise Exception("初始化CTP行情API失败")
            
            print(f"✓ 配置CTP服务器: {md_config['front_addresses'][0]}")
            print(f"✓ 经纪商ID: {md_config['broker_id']}")
            print(f"✓ 投资者ID: {md_config['user_id']}")
            
            # 启动连接
            print("启动CTP行情连接...")
            self.md.start()
            
            # 等待连接和登录
            print("等待连接和登录...")
            max_wait = 30  # 最大等待30秒
            for i in range(max_wait):
                time.sleep(1)
                # 这里可以检查连接状态，但AI_CTP可能没有直接的状态查询方法
                if i >= 5:  # 等待5秒后开始订阅
                    break
                    
            # 订阅行情
            instruments = self.callback.instruments
            print(f"订阅行情: {', '.join(instruments)}")
            
            success = self.md.subscribe(instruments)
            if success:
                print("✓ 订阅成功")
            else:
                print("⚠ 订阅可能失败，请检查合约代码")
            
            # 启用1分钟K线合成（作为备用验证）
            try:
                self.md.enable_kline([1])  # 只启用1分钟K线
                print("✓ 启用1分钟K线合成（备用验证）")
            except Exception as e:
                print(f"⚠ 启用K线合成失败（将使用手动合成）: {e}")
            
            self.running = True
            
            print("=" * 60)
            print("✓ 专业级K线记录器启动完成，开始接收数据...")
            print("📈 使用wtpy标准K线合成，确保数据准确性")
            print("📊 K线合成规则：")
            print("  - 开盘价：时间段内第一个tick价格")
            print("  - 收盘价：时间段内最后一个tick价格") 
            print("  - 最高价：时间段内所有tick的最高价")
            print("  - 最低价：时间段内所有tick的最低价")
            print("  - 成交量：时间段内累计增量")
            print("  - 成交额：时间段内累计增量")
            print("=" * 60)
            
        except Exception as e:
            print(f"✗ 启动K线记录器失败: {e}")
            raise

    def run(self):
        """运行K线记录器"""
        try:
            while self.running:
                time.sleep(5)  # 每5秒打印一次状态
                
                status = self.callback.get_status()
                print(f"📊 运行状态: "
                      f"Tick数={status['tick_count']}, K线数={status['kline_count']}")
                
        except KeyboardInterrupt:
            print("\n收到停止信号，正在停止...")
            self.stop()
            
            final_status = self.callback.get_status()
            print(f"✓ K线记录器已停止，最终统计: "
                  f"处理 {final_status['tick_count']} 个Tick，"
                  f"生成 {final_status['kline_count']} 个标准K线")

    def stop(self):
        """停止K线记录器"""
        self.running = False
        
        # 保存所有未完成的K线
        if self.callback:
            self._save_pending_klines()
        
        if self.md:
            self.md.stop()

    def _save_pending_klines(self):
        """保存所有未完成的K线数据"""
        print("💾 保存未完成的K线数据...")
        
        try:
            saved_count = 0
            
            for instrument_id, current_kline in self.callback.current_klines.items():
                if current_kline is not None and not current_kline.is_completed:
                    # 标记为未完成的K线
                    current_kline.is_completed = False
                    
                    print(f"💾 保存未完成K线: {instrument_id} {current_kline.time_str} "
                          f"OHLC({current_kline.open:.3f},{current_kline.high:.3f},"
                          f"{current_kline.low:.3f},{current_kline.close:.3f}) "
                          f"成交量:{current_kline.volume} Tick数:{current_kline.tick_count} [未闭线]")
                    
                    self.callback._save_kline_to_csv(current_kline)
                    saved_count += 1
            
            if saved_count > 0:
                print(f"💾 已保存 {saved_count} 根未完成的K线")
            else:
                print("💾 无未完成的K线需要保存")
                
        except Exception as e:
            print(f"✗ 保存未完成K线失败: {e}")

def main():
    print("专业级CTP 1分钟K线记录器")
    print("基于wtpy标准的K线合成算法")
    print("=" * 50)
    
    try:
        # 创建K线记录器
        recorder = KlineRecorder()
        
        # 启动记录器
        recorder.start()
        
        # 运行记录器
        recorder.run()
        
    except Exception as e:
        print(f"✗ 程序运行错误: {e}")
    finally:
        print("程序结束")

if __name__ == "__main__":
    main() 